package main

import (
	"ConvertFlow/convertlog"
	"ConvertFlow/protocols/ospf"
	"ConvertFlow/utils"
	"crypto/rsa"
	"fmt"
	"github.com/sirupsen/logrus"
	"log"
	"os"
	"strconv"

	"github.com/spf13/cobra"
)

// OSPFFlags OSPF协议专用的标志
type OSPFFlags struct {
	IfaceName   string
	Filename    string
	ProgressBar bool
	Encryption  bool
}

var ospfFlags = &OSPFFlags{}

func runReceiveByOspf(cmd *cobra.Command, args []string) {
	var privKey *rsa.PrivateKey
	var pubKey *rsa.PublicKey
	var err error
	if ospfFlags.Encryption {
		privKey, pubKey = utils.GenerateKeyPair(4096)
		pubKeyEnc := utils.PublicKeyToBase64(pubKey)
		logrus.Info("Public Key (copy and paste it in receiver): %s", pubKeyEnc)
	}
	size, sender, err := ospf.GetMessageSizeAndSender(ospfFlags.IfaceName)
	if err != nil {
		logrus.Error("ospf.GetMessageSizeAndSender function failed")
		return
	}

	logrus.Infof("Sender: %s, Number of packet wanted: %d", sender, size)
	message, missingPacketsIndexes := ospf.Serve(ospfFlags.IfaceName, size, ospfFlags.ProgressBar)

	//decrypt if Encryption
	if ospfFlags.Encryption {
		message = string(utils.Base64DecryptWithPrivateKey(message, privKey))
	}

	//Print missing packet
	if len(missingPacketsIndexes) > 0 {
		logrus.Warn("Missing packet:")
		for i := 0; i < len(missingPacketsIndexes); i++ {
			fmt.Print(missingPacketsIndexes[i], " ")
		}
		fmt.Println()
	}
	if ospfFlags.Filename != "" {
		f, err := os.Create(ospfFlags.Filename)
		if err != nil {
			logrus.Fatal(err)
		}
		defer f.Close()

		_, err2 := f.WriteString(message)
		if err2 != nil {
			logrus.Fatal(err2)
		}
		logrus.Info("data saved in", ospfFlags.Filename)
	} else {
		logrus.Info("receiver received:", message)
	}
}

func runReceiveByOspfTruncated(cmd *cobra.Command, args []string) {
	var err error
	if ospfFlags.Encryption {
		//privKey, pubKey := utils.GenerateKeyPair(1024)
		_, pubKey := utils.GenerateKeyPair(1024)
		pubKeyEnc := utils.PublicKeyToBase64(pubKey)
		fmt.Println("Public Key (copy and paste it in receiver):")
		fmt.Println(pubKeyEnc)
	}
	size, sender, err := ospf.GetMessageSizeAndSender(ospfFlags.IfaceName)
	if err != nil {
		fmt.Println("ospf.GetMessageSizeAndSender function failed")
		return
	}
	fmt.Println("Sender:", sender, ", Number of packet wanted:", size)
	delay, err := strconv.Atoi(args[0])
	if err != nil {
		panic(err)
	}
	message, missingPacketsIndexes := ospf.ServeTemporary(ospfFlags.IfaceName, size, ospfFlags.ProgressBar, delay)
	//ospfCmd.SendHashedmessage(message, sender) //Integrity check
	//Print missing packet
	if len(missingPacketsIndexes) > 0 {
		fmt.Println("Missing packet:")
		for i := 0; i < len(missingPacketsIndexes); i++ {
			fmt.Print(missingPacketsIndexes[i], " ")
		}
		fmt.Println()
	}

	if ospfFlags.Filename != "" {
		f, err := os.Create(ospfFlags.Filename)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()

		_, err2 := f.WriteString(message)
		if err2 != nil {
			log.Fatal(err2)
		}
		fmt.Println("data saved in", ospfFlags.Filename)
	} else {
		fmt.Println("receiver received:", message)
	}
}

func main() {
	var (
		err error
	)
	//1.load json
	//2.init convertlog
	err = convertlog.InitLogger("client.convertlog", "./", "ERROR")
	if err != nil {
		fmt.Println("InitLogger function failed!")
		return
	}

	var cmdReceive = &cobra.Command{
		Use:   "receive",
		Short: "receive data from ospfCmd packet",
		Long: `receive is for receiving  data from a remote sender.
it uses the ospfCmd protocol.`,
		Args: cobra.MinimumNArgs(0),
		Run:  runReceiveByOspf,
	}

	//cmdReceive flag handling
	cmdReceive.PersistentFlags().StringVarP(&ospfFlags.IfaceName, "listen", "l", "0.0.0.0", "address used for listening ospfCmd packet")
	cmdReceive.PersistentFlags().StringVarP(&ospfFlags.Filename, "filename", "f", "", "Filename where stored the data received")
	cmdReceive.PersistentFlags().BoolVarP(&ospfFlags.ProgressBar, "progress-bar", "p", false, "print progression of the data reception")
	cmdReceive.PersistentFlags().BoolVarP(&ospfFlags.Encryption, "encrypt", "e", false, "use Encryption for data exchange")

	//CMD TRUNCATED
	var cmdTruncated = &cobra.Command{
		Use:   "truncated [delay]",
		Short: "receive data from ospfCmd packet and do not wait indefinitely for all the packet.(indicate the packet missing at the end)",
		Long:  `receive data from ospfCmd packet and do not wait indefinitely for all the packet. Hence you could receive truncated data. If the data isn't fully retrieved  it return the index of the missing packet`,
		Args:  cobra.MinimumNArgs(1),
		Run:   runReceiveByOspfTruncated,
	}

	var rootCmd = &cobra.Command{Use: "receiver"}
	rootCmd.AddCommand(cmdReceive)
	cmdReceive.AddCommand(cmdTruncated)
	rootCmd.Execute()
}
